"""
   matplotlib入门
"""
import matplotlib.pyplot as plt
import numpy as np
from numpy.random import randn
from pandas import Series, DataFrame
import pandas as pd
from datetime import datetime
from mpl_toolkits.basemap import Basemap


def test_01():
    """
    线型图
    :return:
    """
    # s = Series(np.random.randn(10).cumsum(), index=np.arange(0, 100, 10))
    # s.plot()

    df = DataFrame(np.random.randn(10, 4).cumsum(0),
                   columns=['A', 'B', 'C', 'D'],
                   index=np.arange(0, 100, 10))
    df.plot(kind='bar')

    plt.show()


def test_02():
    """
    柱状图
    :return:
    """
    fig, axes = plt.subplots(2, 1)
    data = Series(np.random.rand(16), index=list('abcdefghijklmnop'))
    data.plot(kind='bar', ax=axes[0], color='k', alpha=0.7)
    data.plot(kind='barh', ax=axes[1], color='k', alpha=0.7)
    # plt.show()

    df = DataFrame(np.random.rand(6, 4),
                   index=['one', 'two', 'three', 'four', 'five', 'six'],
                   columns=pd.Index(['A', 'B', 'C', 'D'], name='Genus'))
    df.plot(kind='bar')
    df.plot(kind='bar', stacked=True, alpha=0.5)
    plt.show()


def test_03():
    tips = pd.read_csv('tips.csv')
    party_counts = pd.crosstab(tips.day, tips.size)
    # print(party_counts)
    party_counts = party_counts.ix[:, 2:5]
    # print(party_counts)
    party_pcts = party_counts.div(party_counts.sum(1).astype(float), axis=0)
    print(party_pcts)


def test_04():
    """
    直方图和密度图
    :return:
    """
    comp1 = np.random.normal(0, 1, size=200)
    comp2 = np.random.normal(10, 2, size=200)
    values = Series(np.concatenate([comp1, comp2]))
    values.hist(bins=100, alpha=0.3, color='k', density=True)
    values.plot(kind='kde', style='k--')
    plt.show()


def test_05():
    """
    散布图
    观察两个一维数据序列之间的关系的有效手段
    :return:
    """
    macro = pd.read_csv('macrodata.csv')
    data = macro[['cpi', 'm1', 'tbilrate', 'unemp']]
    trans_data = np.log(data).diff().dropna()
    # print(trans_data[-5:])
    plt.scatter(trans_data['m1'], trans_data['unemp'])
    plt.title('Changes in log %s vs. log %s' % ('m1', 'unemp'))
    pd.plotting.scatter_matrix(trans_data, diagonal='kde', color='k', alpha=0.3)
    plt.show()


def test_06():
    """
    绘制地图：图形化显示海地地震危机数据
    :return:
    """
    data = pd.read_csv('Haiti.csv')
    # print(data)
    # print(data[['INCIDENT DATE', 'LATITUDE', 'LONGITUDE']])[:10]
    # print(data['CATEGORY'][:6])
    # print(data.describe())
    data = data[(data.LATITUDE > 18) & (data.LATITUDE < 20) &
                (data.LONGITUDE > -75) & (data.LONGITUDE < -70)
                & data.CATEGORY.notnull()]
    # print(data)
    all_cats = get_all_categories(data.CATEGORY)
    # 生成器表达式
    english_mapping = dict(get_english(x) for x in all_cats)
    # print(english_mapping['2a'])
    # print(english_mapping['6c'])
    all_codes = get_code(all_cats)
    code_index = pd.Index(np.unique(all_codes))
    dummy_frame = DataFrame(np.zeros((len(data), len(code_index))), index=data.index,
                            columns=code_index)
    # print(dummy_frame.ix[:, :6])
    for row, cat in zip(data.index, data.CATEGORY):
        codes = get_code(to_cat_list(cat))
        dummy_frame.ix[row, codes] = 1
    data = data.join(dummy_frame.add_prefix('category_'))
    # print(data.ix[:, 10:15])
    fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(12, 10))
    fig.subplots_adjust(hspace=0.05, wspace=0.05)

    to_plot = ['2a', 'a', '3c', '7a']

    lllat = 17.25
    urlat = 20.25
    lllon = -75
    urlon = -71

    for code, ax in zip(to_plot, axes.flat):
        m = basic_haiti_map(ax, lllat=lllat, urlat=urlat,
                            lllon=lllon, urlon=urlon)
        cat_data = data[data['category_%s' % code] == 1]

        # 计算地图的投影坐标。
        x, y = m(cat_data.LONGITUDE, cat_data.LATITUDE)

        m.plot(x, y, 'k.', alpha=0.5)
        ax.set_title('%s: %s' % (code, english_mapping[code]))


def basic_haiti_map(ax=None, lllat=17.25, urlat=20.25,
                    lllon=-75, urlon=-71):
    # 创建极球面投影的Basemap实例
    m = Basemap(ax=ax, projection='stere',
                lon_0=(urlon + lllon) / 2,
                lat_0=(urlat + lllat) / 2,
                llcrnrlat=lllat, urcrnrlat=urlat,
                llcrnrlon=lllon, urcrnrlon=urlon,
                resolution='f')
    # 绘制海岸线、州界、国界以及地图边界。
    m.drawcoastlines()
    m.drawstates()
    m.drawcountries()
    return m


def to_cat_list(catstr):
    stripped = (x.strip() for x in catstr.split(','))
    return [x for x in stripped if x]


def get_all_categories(cat_series):
    cat_sets = (set(to_cat_list(x)) for x in cat_series)
    return sorted(set.union(*cat_sets))


def get_english(cat):
    code, names = cat.split('.')
    if '|' in names:
        names = names.split(' | ')[1]
    return code, names.strip()


def get_code(seq):
    return [x.split('.')[0] for x in seq if x]


def main():
    # test_01()
    # test_02()
    # test_03()
    # test_04()
    # test_05()
    test_06()


if __name__ == '__main__':
    main()
